package com.atguigu.gmall.realtime.dwd.db.split.function;

import com.alibaba.fastjson.JSONObject;
import com.atguigu.gmall.realtime.common.bean.TableProcessDwd;
import com.atguigu.gmall.realtime.common.util.JDBCUtil;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.util.Collector;

import java.sql.Connection;
import java.util.*;

public class BaseDbTableProcessFunction extends BroadcastProcessFunction<JSONObject, TableProcessDwd, Tuple2<JSONObject, TableProcessDwd>> {

    private MapStateDescriptor<String, TableProcessDwd> mapStateDescriptor;

    private Map<String, TableProcessDwd> conifMap = new HashMap<>();

    public BaseDbTableProcessFunction(MapStateDescriptor<String, TableProcessDwd> mapStateDescriptor){
        this.mapStateDescriptor = mapStateDescriptor;
    }
    @Override
    public void open(Configuration parameters) throws Exception {
        Connection mysqlConnection = JDBCUtil.getMysqlConnection();
        String sql = "select * from gmall0221_config.table_process_dwd";
        List<TableProcessDwd> tableProcessDwdList = JDBCUtil.queryList(mysqlConnection, sql, TableProcessDwd.class, true);
        for (TableProcessDwd tableProcessDwd : tableProcessDwdList) {
            String sourceTable = tableProcessDwd.getSourceTable();
            String sourceType = tableProcessDwd.getSourceType();
            String key = getKey(sourceTable, sourceType);
            conifMap.put(key, tableProcessDwd);
        }
        JDBCUtil.closeMysqlConnection(mysqlConnection);
    }

    private String getKey(String sourceTable, String sourceType) {
        return sourceTable + ":" + sourceType;
    }

    @Override
    public void processElement(JSONObject jsonObject, BroadcastProcessFunction<JSONObject, TableProcessDwd, Tuple2<JSONObject, TableProcessDwd>>.ReadOnlyContext readOnlyContext, Collector<Tuple2<JSONObject, TableProcessDwd>> collector) throws Exception {
        // {"database":"gmall0221","xid":925472,"data":{"order_status":"1006","create_time":"2024-07-23 19:13:03","id":19039,"order_id":9726},"commit":true,"type":"insert","table":"order_status_log","ts":1722424383}
        String table = jsonObject.getString("table");
        String type = jsonObject.getString("type");
        String key = getKey(table, type);
        ReadOnlyBroadcastState<String, TableProcessDwd> broadcastState = readOnlyContext.getBroadcastState(mapStateDescriptor);
        TableProcessDwd tableProcessDwd = null;
        if ((tableProcessDwd = broadcastState.get(key)) != null
        || (tableProcessDwd = conifMap.get(key)) != null){
            JSONObject dataObj = jsonObject.getJSONObject("data");
            String sinkColumns = tableProcessDwd.getSinkColumns();
            deleteNotNeedCol(dataObj, sinkColumns);
            Long ts = jsonObject.getLong("ts");
            dataObj.put("ts", ts);
            collector.collect(Tuple2.of(dataObj, tableProcessDwd));
        }
    }

    private void deleteNotNeedCol(JSONObject dataObj, String sinkColumns) {
        List<String> columnList = Arrays.asList(sinkColumns.split(","));
        Set<Map.Entry<String, Object>> entrySet = dataObj.entrySet();
        entrySet.removeIf(entry -> columnList.contains(entry.getKey()));
    }

    @Override
    public void processBroadcastElement(TableProcessDwd tableProcessDwd, BroadcastProcessFunction<JSONObject, TableProcessDwd, Tuple2<JSONObject, TableProcessDwd>>.Context context, Collector<Tuple2<JSONObject, TableProcessDwd>> collector) throws Exception {
        String op = tableProcessDwd.getOp();
        BroadcastState<String, TableProcessDwd> broadcastState = context.getBroadcastState(mapStateDescriptor);
        String sourceTable = tableProcessDwd.getSourceTable();
        String sourceType = tableProcessDwd.getSourceType();
        String key = getKey(sourceTable, sourceType);
        if ("d".equals(op)){
            broadcastState.remove(key);
            conifMap.remove(key);
        } else {
            broadcastState.put(key, tableProcessDwd);
            conifMap.put(key, tableProcessDwd);
        }
    }

}


